'''
Created on Apr 30, 2012
@author: alebalbin
'''

import glob
import os
import sys
import subprocess
from collections import defaultdict

JOB_ERROR=1
JOB_SUCCESS=0

def read_files_folder(folderpath,ext,minsp,maxsp):
    ''' '''
    # Read files in folder
    myfiles=defaultdict()
    for infile in glob.glob( os.path.join(folderpath, '*'+ext) ):
        filename=infile.split('/')[-1]
        sp=filename[0:8]
        spnum=int(sp[6:8])
        if (spnum > minsp) and (spnum < maxsp):
            myfiles[sp]=os.path.join(folderpath,filename)
    return myfiles

def make_symb_link(sample_dict,dirbase):
    '''
    '''    
    for sp,f in sample_dict.iteritems():
        fs=os.path.join(dirbase,sp+'.CEL')
        args=['ln','-s',f,fs]
        print args
        r=subprocess.call(args)
        
    return JOB_SUCCESS


def preprocessing_celfiles(dirbase,celfiledir,output_raw):
    '''
    java -Xmx2G -jar ./celConverter/CelFileConverter.jar -m 
    ./celConverter/Snp6FeatureMappings.csv 
    -c ./cdf/GenomeWideSNP_6.Full.cdf -s ./celfile/ -t ./outdir/raw/
    '''
    os.chdir(dirbase)
    rawdir = os.path.join(dirbase,'outdir/'+output_raw)
    if not os.path.exists(rawdir):
        os.mkdir(rawdir)
        
    args=['java','-Xmx2G','-jar', 
          os.path.join(dirbase,'celConverter/CelFileConverter.jar'),'-m', 
          os.path.join(dirbase,'celConverter/Snp6FeatureMappings.csv'), '-c', 
          os.path.join(dirbase,'cdf/GenomeWideSNP_6.Full.cdf'), '-s',
          celfiledir, '-t', 
          os.path.join(dirbase,'outdir/'+output_raw+'/')]
    print args
    r = subprocess.call(args)
    
    if r!=0:
        return JOB_ERROR
    return JOB_SUCCESS

def normalise_and_segmentation(dirbase,dirbase_outdir,output_raw):
    '''
    ./run_preprocessing.sh ./Matlab_Compiler_Runtime/v710 Beer_001_\(GenomeWideSNP_6\).feature_intensity 
    ./info/ ./outdir/raw/ ./outdir/Beer_001/output/ ./outdir/Beer_001/ PRIMARY 0.1
    catch 3 final numbers in the last line.
    '''
    os.chdir(dirbase)
    rawdata = os.path.join(dirbase_outdir,output_raw)
    rawfiles = os.listdir(rawdata)
    analysis_done=[]
    #catch 3 final numbers in the last line.]
    
    for file1 in rawfiles:
        print file1
        file1=file1.strip('/')
        print file1
        sample_name=file1[0:8]
        dirsample=os.path.join(dirbase_outdir,sample_name)
        dirsp_output=os.path.join(dirbase_outdir,sample_name+'/output')
        if not os.path.exists(dirsample):
            os.mkdir(dirsample)
        if not os.path.exists(dirsp_output):
            os.mkdir(dirsp_output)
        
        #normalization and ploidy estimation
        args = [os.path.join(dirbase,'run_preprocessing.sh'), 
                os.path.join(dirbase,'Matlab_Compiler_Runtime/v710'), file1, 
                os.path.join(dirbase,'info/'),
                os.path.join(dirbase_outdir,output_raw+'/'), 
                dirsp_output+'/', 
                dirsample+'/', 
                'PRIMARY', '0.1']
        
        args=",".join(args).replace(',', ' ')
        # segmentation
        print args
        r = subprocess.Popen(args,stdout=subprocess.PIPE,shell=True)
        t=r.communicate()
        param_list=t[0].split('\n')[-2].split()
        retcode=segmentation(dirbase,param_list,sample_name,file1)
        
        if retcode!=0:
            analysis_done.append((sample_name,JOB_ERROR))
        
    return analysis_done

def segmentation(dirbase,param_list,sample,file_name):
    '''
    ./run_HMM.sh ./Matlab_Compiler_Runtime/v710/ Beer_001_\(GenomeWideSNP_6\).feature_intensity 
    ./info/ ./outdir/output/ ./outdir/ 10 0.1129 1.7418 0.40442
    '''
    args=[os.path.join(dirbase,'run_HMM.sh'),
          os.path.join(dirbase,'Matlab_Compiler_Runtime/v710'),
          os.path.join(dirbase,file_name),
          os.path.join(dirbase,'info/'),
          os.path.join(dirbase,'outdir/'+sample+'/output/'),
          os.path.join(dirbase,'outdir/'+sample+'/'),
          '10',param_list[0],param_list[1],param_list[2]]

    args=",".join(args).replace(',', ' ')
    print args
    r = subprocess.call(args,shell=True)
    
    if r!=0:
        return JOB_ERROR
    return JOB_SUCCESS
    
    

if __name__ == '__main__':
    celfolder='/exds/users/oabalbin/projects/nsclc_kras_dep/trunk/Data/SNP6_arrays/'
    ext='CEL'
    dirbase='/exds/users/oabalbin/sw/picnic'
    dirbase_outdir='/exds/users/oabalbin/sw/picnic/outdir'
    celfiledir=os.path.join(dirbase,'celfile')
    #output_raw='raw'
    samples_dict=read_files_folder(celfolder,ext, int(sys.argv[1]),int(sys.argv[2]) )
    print samples_dict
    output_raw='raw_%s_%s'%(sys.argv[1],sys.argv[2])
    celfiledir=os.path.join(celfiledir,output_raw)
    if not os.path.exists(celfiledir):
        os.mkdir(celfiledir)
    make_symb_link(samples_dict,celfiledir)
    # Preprocess files
    r=preprocessing_celfiles(dirbase,celfiledir,output_raw)
    retcode=JOB_ERROR
    if r !=0:
        pass
    else:
        analysis_done=normalise_and_segmentation(dirbase,dirbase_outdir,output_raw)
        print analysis_done
        retcode=JOB_SUCCESS
    print retcode
    
    
    

